package com.goldfish666.houghrecognition.util;

import lombok.extern.slf4j.Slf4j;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.*;

/**
 * 预处理程序
 * 包含灰度化、二值化、骨架化、截取仅含内容的图片
 */
@Slf4j
public class HoughImageUtil {


    public static BufferedImage read(String path) throws IOException {
        return ImageIO.read(new File(path));
    }

    public static void draw(BufferedImage bufferedImage, String path, String name) throws IOException {
        log.info("打印图片{}", name);
        ImageIO.write(bufferedImage, "jpg", new File(path + name + ".jpg"));
    }

    // 灰度化
    public static BufferedImage getGrayImage(BufferedImage image) throws IOException {

        BufferedImage grayBufferedImage = new BufferedImage(image.getWidth(), image.getHeight(), image.getType());

        for (int i = 0; i < image.getWidth(); i++) {
            for (int j = 0; j < image.getHeight(); j++) {
                int value = image.getRGB(i, j);
                int gray = (38 * getRed(value) + 75 * getGreen(value) + 15 * getBlue(value)) >> 7;
                int color = new Color(gray, gray, gray).getRGB();
                grayBufferedImage.setRGB(i, j, color);
            }
        }
        return grayBufferedImage;
    }

    public static int getRed(int point) {
        return (point >> 16) & 0xff;
    }

    public static int getGreen(int point) {
        return (point >> 8) & 0xff;
    }

    public static int getBlue(int point) {
        return point & 0xff;
    }

    public static int getGray(int point) {
        return point & 0xff;
    }

    // 二值化
    public static BufferedImage getBinaryImage(BufferedImage bufferedImage) {
        int w = bufferedImage.getWidth();
        int h = bufferedImage.getHeight();
        int black = new Color(0, 0, 0).getRGB();
        int white = new Color(255, 255, 255).getRGB();
        BufferedImage bi = new BufferedImage(w, h, bufferedImage.getType());
        int T = getThreshold(bufferedImage);
        for (int x = 0; x < w; x++) {
            for (int y = 0; y < h; y++) {
                int pixel = bufferedImage.getRGB(x, y);
                int gray = getGray(pixel);

                if (gray < T) {
                    bi.setRGB(x, y, black);
                } else {
                    bi.setRGB(x, y, white);
                }
            }
        }
        return bi;
    }

    // 获取二值化的阈值
    public static int getThreshold(BufferedImage bufferedImage) {
        Map<Integer, Double> map = new LinkedHashMap<>(256);
        int width = bufferedImage.getWidth();
        int height = bufferedImage.getHeight();

        for (int i = 0; i < 256; i++) {
            List<Integer> foreground = new ArrayList<>();
            List<Integer> background = new ArrayList<>();
            for (int x = 0; x < width; x++) {
                for (int y = 0; y < height; y++) {
                    int pixel = bufferedImage.getRGB(x, y);
                    int gray = getGray(pixel);

                    if (gray < i) {
                        foreground.add(gray);
                    } else {
                        background.add(gray);
                    }
                }
            }

            int foregroundCount = foreground.size();
            int backgroundCount = background.size();
            if (foregroundCount * backgroundCount != 0) {
                int total = width * height;

                double foregroundAverage = getAverageOfList(foreground);
                double backgroundAverage = getAverageOfList(background);
                double entiretyAverage = foregroundCount * 1.0 / total * foregroundAverage + backgroundCount * 1.0 / total * backgroundAverage;

                double variance = foregroundCount * 1.0 / total * Math.pow(entiretyAverage - foregroundAverage, 2) + backgroundCount * 1.0 / total * Math.pow(entiretyAverage - backgroundAverage, 2);
                map.put(i, variance);
            }
        }
        return map.entrySet().stream().sorted(Comparator.comparingInt(e -> -(int) (e.getValue() * 1000))).findFirst().get().getKey();
    }

    public static double getAverageOfList(List<Integer> list) {
        int sum = 0;
        for (Integer integer : list) {
            sum += integer;
        }
        return sum * 1.0 / list.size();
    }

    // 骨架化
    public static BufferedImage skeletonize(BufferedImage bufferedImage) {
        int w = bufferedImage.getWidth();
        int h = bufferedImage.getHeight();
        int white = new Color(255, 255, 255).getRGB();
        BufferedImage result = new BufferedImage(w, h, bufferedImage.getType());
        int[][] number = new int[h][w];
        int T = getThreshold(bufferedImage);
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                number[i][j] = getGray(bufferedImage.getRGB(j, i)) < T ? 1 : 0;
                result.setRGB(j, i, bufferedImage.getRGB(j, i));
            }
        }

        List<Pair> toClear = new ArrayList<>();
        do {
            toClear.clear();
            for (int i = 0; i < h; i++) {
                for (int j = 0; j < w; j++) {
                    if (number[i][j] == 1) {
                        Map<String, Integer> map0 = getPMap(i, j, number);
                        if (aroundForegroundNumberTest(map0) && backgroundForegroundPatternTest(map0) && rightBackgroundTest(map0) && bottomBackgroundTest(map0)) {
                            toClear.add(new Pair(i, j));
                        }
                    }
                }
            }

            for (Pair p : toClear) {
                number[p.getX()][p.getY()] = 0;
            }

            for (int i = 0; i < h; i++) {
                for (int j = 0; j < w; j++) {
                    if (number[i][j] == 1) {
                        Map<String, Integer> map0 = getPMap(i, j, number);
                        if (aroundForegroundNumberTest(map0) && backgroundForegroundPatternTest(map0) && topTriangleBackgroundTest(map0) && leftTriangleBackgroundTest(map0)) {
                            toClear.add(new Pair(i, j));
                        }
                    }

                }
            }

            for (Pair p : toClear) {
                number[p.getX()][p.getY()] = 0;
            }

        } while (toClear.size() > 0);

        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                if (number[i][j] == 0) {
                    result.setRGB(j, i, white);
                }
            }
        }

        return result;
    }

    public static boolean leftTriangleBackgroundTest(Map<String, Integer> map) {
        return map.get("p8") * map.get("p2") * map.get("p6") == 0;
    }

    public static boolean topTriangleBackgroundTest(Map<String, Integer> map) {
        return map.get("p8") * map.get("p4") * map.get("p2") == 0;
    }

    public static boolean bottomBackgroundTest(Map<String, Integer> map) {
        return map.get("p8") * map.get("p4") * map.get("p6") == 0;
    }

    public static boolean rightBackgroundTest(Map<String, Integer> map) {
        return map.get("p2") * map.get("p4") * map.get("p6") == 0;
    }

    public static boolean backgroundForegroundPatternTest(Map<String, Integer> map) {
        StringBuilder s = new StringBuilder();
        for (int i = 2; i < map.size() + 1; i++) {
            s.append(map.get("p" + i));
        }
        s.append(map.get("p2"));
        String sx = s.toString();
        return sx.contains("01") && sx.indexOf("01") == sx.lastIndexOf("01");
    }

    public static boolean aroundForegroundNumberTest(Map<String, Integer> map) {
        int sum = -map.get("p1");
        for (Map.Entry<String,Integer> entry: map.entrySet()){
            sum += entry.getValue();
        }
        return 2 <= sum && sum <= 6;
    }

    public static Map<String, Integer> getPMap(int x, int y, int[][] number) {
        int h = number.length;
        int w = number[0].length;
        int p1 = number[x][y];
        int p2 = x == 0 ? 0 : number[x - 1][y];
        int p3 = y == w - 1 || x == 0 ? 0 : number[x - 1][y + 1];
        int p4 = y == w - 1 ? 0 : number[x][y + 1];
        int p5 = y == w - 1 || x == h - 1 ? 0 : number[x + 1][y + 1];
        int p6 = x == h - 1 ? 0 : number[x + 1][y];
        int p7 = y == 0 || x == h - 1 ? 0 : number[x + 1][y - 1];
        int p8 = y == 0 ? 0 : number[x][y - 1];
        int p9 = y == 0 || x == 0 ? 0 : number[x - 1][y - 1];
        Map<String, Integer> result = new HashMap<>();
        result.put("p1", p1);
        result.put("p2", p2);
        result.put("p3", p3);
        result.put("p4", p4);
        result.put("p5", p5);
        result.put("p6", p6);
        result.put("p7", p7);
        result.put("p8", p8);
        result.put("p9", p9);
        return result;
    }

    static class Pair {
        private final int x;
        private final int y;

        public Pair(int x, int y) {
            this.x = x;
            this.y = y;
        }

        public int getX() {
            return x;
        }

        public int getY() {
            return y;
        }
    }

    private static Map<Integer, Integer> getXCountMap(BufferedImage bufferedImage) {
        Map<Integer, Integer> map = new HashMap<>(bufferedImage.getWidth());
        for (int i = 0; i < bufferedImage.getWidth(); i++) {
            int count = 0;
            for (int j = 0; j < bufferedImage.getHeight(); j++) {
                if (getBlue(bufferedImage.getRGB(i, j)) == 0) {
                    count++;
                    break;
                }
            }
            map.put(i, count);
        }
        return map;
    }

    public static int getStartX(BufferedImage bufferedImage) {
        Map<Integer, Integer> map = getXCountMap(bufferedImage);
        int left = map.entrySet().stream().sorted(Comparator.comparingInt(Map.Entry::getKey)).filter(e -> e.getValue() != 0).findFirst().get().getKey();
        return Math.max(left - 1, 0);
    }

    public static int getEndX(BufferedImage bufferedImage) {
        Map<Integer, Integer> map = getXCountMap(bufferedImage);
        int right = map.entrySet().stream().sorted((o1, o2) -> o2.getKey() - o1.getKey()).filter(e -> e.getValue() != 0).findFirst().get().getKey();
        return Math.min(right + 1, bufferedImage.getTileWidth() - 1);
    }

    private static Map<Integer, Integer> getYCountMap(BufferedImage bufferedImage) {
        Map<Integer, Integer> map = new HashMap<>(bufferedImage.getWidth());
        for (int i = 0; i < bufferedImage.getHeight(); i++) {
            int count = 0;
            for (int j = 0; j < bufferedImage.getWidth(); j++) {
                if (getGray(bufferedImage.getRGB(j, i)) == 0) {
                    count++;
                    break;
                }
            }
            map.put(i, count);
        }
        return map;
    }

    public static int getStartY(BufferedImage bufferedImage) {
        Map<Integer, Integer> map = getYCountMap(bufferedImage);
        int top = map.entrySet().stream().sorted(Comparator.comparingInt(Map.Entry::getKey)).filter(e -> e.getValue() != 0).findFirst().get().getKey();
        return Math.max(top - 1, 0);
    }

    public static int getEndY(BufferedImage bufferedImage) {
        Map<Integer, Integer> map = getYCountMap(bufferedImage);
        int bottom = map.entrySet().stream().sorted((o1, o2) -> o2.getKey() - o1.getKey()).filter(e -> e.getValue() != 0).findFirst().get().getKey();
        return Math.min(bottom + 1, bufferedImage.getHeight() - 1);
    }

    public static BufferedImage getContentImage(BufferedImage bufferedImage) {
        int startX = getStartX(bufferedImage);
        int endX = getEndX(bufferedImage);
        int startY = getStartY(bufferedImage);
        int endY = getEndY(bufferedImage);
        BufferedImage result = new BufferedImage(endX - startX + 1, endY - startY + 1, bufferedImage.getType());
        for (int i = startX; i <= endX; i++) {
            for (int j = startY; j <= endY; j++) {
                result.setRGB(i - startX, j - startY, bufferedImage.getRGB(i, j));
            }
        }
        return result;
    }

    static class ImageInfo {

        private final int width;
        private final int height;

        private final int T;

        public ImageInfo(int width, int height, int t) {

            this.width = width;
            this.height = height;

            this.T = t;
        }

        public int getWidth() {
            return width;
        }

        public int getHeight() {
            return height;
        }

        public int getT() {
            return T;
        }
    }

    // 针对已骨架化的图片
    public static ImageInfo getImageInfo(BufferedImage bufferedImage) {
        return new ImageInfo( bufferedImage.getWidth(), bufferedImage.getHeight(), getThreshold(bufferedImage));
    }


}
